'''
Created on Oct 8, 2009

@author: mkiyer
'''
from veggie.app.alignment.bowtie import Bowtie
from veggie.sample.sampledb2 import get_sampledb
import sys
import logging
import os
import glob
import re

def find_sequence_files(flowcell, lane, sequence_path,
                        prefix_fmt="%s_%s",
                        seq_suffix="_best_sequence.txt",
                        bowtie_quals_suffix='_best_bowtiequals.txt'):

    analysis_id = prefix_fmt % (flowcell, lane)
    # build a string to use to search for sequence files
    seq_search_string = ''.join([analysis_id, '*', seq_suffix])
    
    # glob to get all the sequence file matches
    fastq_files = glob.glob(os.path.join(sequence_path, seq_search_string))
    if len(fastq_files) == 0:
        logging.error('No sequence files for flowcell %s lane %s' % (flowcell, lane))
        return

    # regexp for getting beginning of filename to use as prefix
    # in search for bowtie quals file
    prefix_re = re.compile(r'(.*)' + seq_suffix)
    
    for fastq_file in fastq_files:
        m = prefix_re.match(os.path.basename(fastq_file))
        if not m:
            logging.critical("regexp match failed on file %s" % fastq_file)
            continue
        # get the quals param from the bowtie file
        fastq_prefix = m.group(1)
        quals_file = os.path.join(sequence_path, fastq_prefix + bowtie_quals_suffix)            
        if not os.path.exists(quals_file):
            logging.critical("quality score file %s not found... skipping %s" %
                             (quals_file, fastq_file))
            continue
        quals_fhd = open(quals_file)
        quals_param = quals_fhd.readline()
        quals_fhd.close()
        # return tuple of fastq_file and quals_param
        yield (fastq_prefix, fastq_file, quals_param)


def virus_align(analysis_id, fastq_file, quals_param,
                bowtie_cwd):    
    # construct output file names
    aligned_file = '%s_aligned.txt' % (analysis_id)
    unaligned_file = '%s_unaligned.fq' % (analysis_id)
    maxmms_file = '%s_maxmms.fq' % (analysis_id)
    virus_file = '%s_viruses.txt' % (analysis_id)

    # skip alignments that already exist
    logging.debug("file: %s quals: %s output_file: %s" % (fastq_file, quals_param, virus_file))
    if os.path.exists(os.path.join(bowtie_cwd, virus_file)):
        logging.debug("output file %s exists, skipping..." % virus_file)
        return

    # setup bowtie
    hg18_bowtie_options = {'-p': 8,
                           '--' + quals_param: None,
                           '--un': unaligned_file,
                           '-m': 32,
                           '--max': maxmms_file}    
    # create the app controller
    bowtie_app = Bowtie(cwd=bowtie_cwd, options=hg18_bowtie_options)        
    # align to human genome
    bowtie_app.run_single_best(fastq_file, aligned_file, genome="hg18")

    # align the reads that could not be aligned to the human genome to viral genomes        
    virus_bowtie_options = {'-p': 8,
                            '--' + quals_param: None}
    virus_bowtie_app = Bowtie(cwd=bowtie_cwd, options=virus_bowtie_options)
    virus_bowtie_app.run_single_best(unaligned_file, virus_file, genome="viruses_cds")
    
    # delete intermediate files
    logging.debug("deleting intermediate files...")
    os.remove(aligned_file)
    os.remove(maxmms_file)
    #os.remove(unaligned_file)
    logging.debug("finished %s" % (analysis_id))

def virus_search(fclanes, sequence_path, bowtie_cwd):
    for flowcell, lane in fclanes:
        for res in find_sequence_files(flowcell, lane, sequence_path):
            fastq_prefix, fastq_file, quals_param = res
            virus_align(fastq_prefix, fastq_file, quals_param, bowtie_cwd) 

def get_ribominus_analyses():
    # query sampleDB for flowcell/lanes
    fclanes = get_sampledb().get_analyses({'app_type': 'RiboMinus'})
    return fclanes

def get_paired_transcriptome_analyses():
    return (('302LVAAXX', 2), ('42B3YAAXX', 6))

def get_transcriptome_analyses():
    # query sampleDB for flowcell/lanes
    fclanes = get_sampledb().get_analyses({'app_type': 'Transcriptome'})
    return fclanes

if __name__ == '__main__':
    logging.basicConfig(format='%(asctime)s %(levelname)s %(message)s',
                        level=logging.DEBUG)
    sequence_path = sys.argv[1]
    dest_path = sys.argv[2]

    # make directory for results
    bowtie_cwd = dest_path
    if not os.path.exists(bowtie_cwd):
        logging.debug('Creating bowtie results path %s' % bowtie_cwd)
        os.makedirs(bowtie_cwd)    

    # start search
    #virus_search(get_ribominus_analyses(), sequence_path, dest_path)
    #virus_search(get_paired_transcriptome_analyses(), sequence_path, dest_path)
    virus_search(get_transcriptome_analyses(), sequence_path, dest_path)
    #bowtie_best_lanes(sequence_path, dest_path)
    
    
